from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.metrics import roc_auc_score, classification_report, roc_curve, auc
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.model_selection import GridSearchCV
import joblib
from xgboost import XGBClassifier

plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体为黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像时负号 '-' 显示为方块的问题

# 读取数据
data = pd.read_csv("../../data/raw/train.csv")

# 二分类特征进行独热编码
binary_cols = ['Gender', 'Over18', 'OverTime']
data = pd.get_dummies(data, columns=binary_cols, drop_first=True)

# 多分类特征进行更细致的大分类并赋予相应权重

# BusinessTravel 特征
# 频繁出差可能对工作和生活平衡影响较大，权重可适当提高
travel_mapping = {'Non-Travel': 0, 'Travel_Rarely': 1, 'Travel_Frequently': 3}
data['BusinessTravel'] = data['BusinessTravel'].map(travel_mapping)

# Department 特征
# 研发部门可能对专业技能要求高，销售部门注重业绩，人力资源部门侧重协调管理，可据此赋权
department_mapping = {'Research & Development': 0, 'Sales': 2, 'Human Resources': 1}
data['Department'] = data['Department'].map(department_mapping)

# EducationField 特征
# 生命科学和医学可能专业性较强，与研发等岗位相关性高，可赋予较高权重
education_mapping = {'Other': 0, 'Human Resources': 1, 'Marketing': 2, 'Technical Degree': 3, 'Medical': 4, 'Life Sciences': 4}
data['EducationField'] = data['EducationField'].map(education_mapping)

# JobRole 特征，分为管理岗和非管理岗
management_roles = ['Manufacturing Director', 'Research Director', 'Manager']
jobrole_mapping = {role: 1 if role in management_roles else 0 for role in data['JobRole'].unique()}
data['JobRole'] = data['JobRole'].map(jobrole_mapping)

# MaritalStatus 特征
# 已婚人士可能更稳定，单身人士可能更有精力投入新挑战，离异人士可能处于调整期，可按此赋权
marital_mapping = {'Single': 1, 'Married': 2, 'Divorced': 0}
data['MaritalStatus'] = data['MaritalStatus'].map(marital_mapping)

# 选取特征列  12 + 9列
# 上边通过卡方检验 选出了 12列
# 下边通过 随机森林 + 置换重要性选出了 9列
X = data[
    [
        # 商务差旅频率（分类）
        "BusinessTravel",

        # 所在部门（分类）
        "Department",

        # 教育领域（分类）
        "EducationField",

        # 工作环境满意度评分（1-4）
        "EnvironmentSatisfaction",

        # 工作投入度评分（1-4）
        "JobInvolvement",

        # 职位级别（数值）
        "JobLevel",

        # 工作角色（分类）
        "JobRole",

        # 工作满意度评分（1-4）
        "JobSatisfaction",

        # 婚姻状况（分类）
        "MaritalStatus",

        # 是否加班（布尔）
        "OverTime_Yes",

        # 股票期权等级（0-3）
        "StockOptionLevel",

        # 工作生活平衡评分（1-4）
        "WorkLifeBalance",

        # 年龄（数值）
        "Age",

        # 家到公司的距离（公里）
        "DistanceFromHome",

        # 上次晋升至今的年数
        "YearsSinceLastPromotion",

        # 在当前职位的年数
        "YearsInCurrentRole",

        # 曾工作过的公司数量
        "NumCompaniesWorked",

        # 月收入（数值）
        "MonthlyIncome",

        # 与当前经理共事的年数
        "YearsWithCurrManager",

        # 薪资增长百分比
        "PercentSalaryHike",

        # 总工作年限
        "TotalWorkingYears",
    ]
]
y = data["Attrition"]
# 对X 标准化
scale = StandardScaler()
scale.fit(X)
X = scale.transform(X)
joblib.dump(scale, "scale_model.pkl")

# 划分数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=25)
# 使用XGBClassifier 进行模型训练
es_xgb = LogisticRegression(random_state=25, max_iter=1000)  # max_iter 设置较大值以确保收敛

es_xgb.fit(X_train, y_train)
y_pred = es_xgb.predict_proba(X_test)[:, 1]

# 打印auc
print("AUC:", roc_auc_score(y_test, y_pred))  # AUC: 0.9005011257171909

# 保存模型
joblib.dump(es_xgb, "xgb_model.pkl")
